
import os
os.environ['DISABLE_ARGUMENTS_CHECKER'] = "Yes"

from rqalpha.apis import *


from funcat import *

# 在这个方法中编写任何的初始化逻辑。context对象将会在你的算法策略的任何方法之间做传递。
def init(context):
    context.s1 = context.config.mod.sys_analyser.benchmark

    # 初始化画图类
    # plot_init(context.s1)
    #
    context.is_got_bars = False
    context.long_days = None
    context.long_holds = {}
    context.short_days = None
    context.short_holds = {}
    context.dkx = 0
    context.madkx = 0
    context.freq = context.config.base.frequency
    context.rdb = RQAlphaDataBackend(data_proxy=Environment.get_instance().data_proxy)

# 你选择的证券的数据更新将会触发此段逻辑，例如日或分钟历史数据切片或者是实时数据切片更新
def handle_bar(context, bar_dict):
    # 开始编写你的主要的算法逻辑
    # bar_dict[order_book_id] 可以拿到某个证券的bar信息
    # context.portfolio 可以拿到现在的投资组合状态信息
    # 停牌直接返回
    if is_suspended(context.s1):
      return
    if not context.is_got_bars:
      bars = history_bars(context.s1, 20+1, '1d', adjust_type='none')
      if len(bars) < 21:
        return
      context.is_got_bars = True
    with funcat_execution_context(date=context.now,
                         order_book_id=context.s1,
                         data_backend=context.rdb, 
                         freq=context.freq) as fec:
        dkx,madkx = DKX()
        # plot_line("#0dkx", dkx.value)
        # plot_line("#0madkx", madkx.value)
        # 用上个周期的多空线判断
        if context.dkx > context.madkx:
          # plot_icon(2)
          # 有空仓，先平仓
          if context.s1 in context.short_holds:
            # plot_order("BUY_CLOSE", 50)
            target_shares = context.short_holds[context.s1]
            ret = order_shares(context.s1, target_shares, price=OPEN.value)
            if ret and ret.status == ORDER_STATUS.FILLED:
                del context.short_holds[context.s1]
            print("BUY_CLOSE", context.dkx, context.madkx, dkx, madkx, OPEN)
          if context.s1 not in context.long_holds:
            # 计算箱体，向上突破就做多
            tlrb = BOX_FIND(-1, context.long_days)
            if tlrb and HIGH > tlrb[0]:
              # plot_order("BUY_OPEN", 100)
              price = max(tlrb[0].value + 0.01, OPEN.value)
              ret = order_target_percent(context.s1, 1, price=price)
              if ret and ret.filled_quantity > 0:
                context.long_holds[context.s1] = context.now.date()
              print("BUY_OPEN", context.dkx, context.madkx, dkx, madkx, HIGH, tlrb, price)
              pass
        elif context.dkx < context.madkx:
          # plot_icon(3)
          # 有多仓，先清仓
          if context.s1 in context.long_holds:
            # plot_order("SELL_CLOSE", 100)
            ret = order_target_percent(context.s1, 0, price=OPEN.value)
            if ret and ret.status == ORDER_STATUS.FILLED:
                del context.long_holds[context.s1]
            print("SELL_CLOSE", context.dkx, context.madkx, dkx, madkx, OPEN)
          # 是否支持做空
          if not context.config.mod.sys_accounts.validate_stock_position:
            if context.s1 not in context.short_holds:
              # 计算箱体，向下突破就做空
              blrt = BOX_FIND(1, context.short_days)
              if blrt and LOW < blrt[0]:
                # plot_order("SELL_OPEN", 50)
                target_shares = int(context.portfolio.total_value * 0.5 / blrt[0].value / 100) * 100
                price = min(blrt[0].value - 0.01, OPEN.value)
                ret = order_shares(context.s1, -target_shares, price=price)
                if ret and ret.filled_quantity > 0:
                    context.short_holds[context.s1] = target_shares
                print("SELL_OPEN", context.dkx, context.madkx, dkx, madkx, LOW, blrt, price)
                pass
        # 保存日期
        if dkx > madkx:
          context.short_days = None
          if context.long_days is None:
            print("Long:", context.dkx, context.madkx, dkx, madkx)
            context.long_days = 3
          else:
            context.long_days += 1
        elif dkx < madkx:
          context.long_days = None
          if context.short_days is None:
            print("Short:", context.dkx, context.madkx, dkx, madkx)
            context.short_days = 3
          else:
            context.short_days += 1
        # 保存上一个多空线
        context.dkx = dkx
        context.madkx = madkx
    # 持仓
    if len(context.long_holds) > 0:
      print("long:", context.long_holds)
    if len(context.short_holds) > 0:
      print("short:", context.short_holds)


__config__ = {
  "base": {
    "strategy_name": "dkx_after_box",
    "start_date": "2010-01-01",
    "end_date": "2022-09-01",
    "accounts": {
        "stock": 1000000
    }
  },
  "extra": {
    "log_level": "warning",
  },
  "mod": {
    "sys_accounts": {
      "enabled": True,
      "validate_stock_position": True
    },
    "sys_analyser": {
      "benchmark": "000333.XSHE",
      "enabled": True,
      "plot": True
    },
    "sys_simulation": {
      "enabled": True,
      "signal": True
    },
    "stu_plot": {
      "enabled": True,
      "to_sql": True
    }
  }
}


if __name__ == '__main__':
    from rqalpha import run_file
    import json
    import time
    # import cProfile
    # pr = cProfile.Profile()
    # pr.enable()
    # 您可以指定您要传递的参数
    ret = run_file(__file__, config=__config__)
    print(ret)
    pd.set_option('display.max_rows', None)
    print(ret['sys_analyser']['trades'])
    # pr.disable()
    # pr.print_stats('cumtime')
    exit()
    # ["000166.XSHE", "000333.XSHE"]
    # 批量回测
    # code_list = \
    # ['000001.XSHE', '000002.XSHE', '000063.XSHE', '000066.XSHE', '000069.XSHE', '000100.XSHE', '000157.XSHE', '000166.XSHE', '000301.XSHE', '000333.XSHE', '000338.XSHE', '000408.XSHE', '000425.XSHE', '000538.XSHE', '000568.XSHE', 
    # '000596.XSHE', '000625.XSHE', '000651.XSHE', '000661.XSHE', '000703.XSHE', '000708.XSHE', '000725.XSHE', '000768.XSHE', '000776.XSHE', '000786.XSHE', '000792.XSHE', '000800.XSHE', '000858.XSHE', '000876.XSHE', '000877.XSHE', 
    # '000895.XSHE', '000938.XSHE', '000963.XSHE', '000977.XSHE', '001289.XSHE', '001979.XSHE', '002001.XSHE', '002007.XSHE', '002008.XSHE', '002027.XSHE', '002032.XSHE', '002049.XSHE', '002050.XSHE', '002064.XSHE', '002074.XSHE', 
    # '002120.XSHE', '002129.XSHE', '002142.XSHE', '002179.XSHE', '002202.XSHE', '002230.XSHE', '002236.XSHE', '002241.XSHE', '002252.XSHE', '002271.XSHE', '002304.XSHE', '002311.XSHE', '002352.XSHE', '002371.XSHE', '002410.XSHE', 
    # '002414.XSHE', '002415.XSHE', '002459.XSHE', '002460.XSHE', '002466.XSHE', '002475.XSHE', '002493.XSHE', '002555.XSHE', '002568.XSHE', '002594.XSHE', '002600.XSHE', '002601.XSHE', '002602.XSHE', '002607.XSHE', '002648.XSHE', 
    # '002709.XSHE', '002714.XSHE', '002736.XSHE', '002791.XSHE', '002812.XSHE', '002821.XSHE', '002841.XSHE', '002916.XSHE', '002920.XSHE', '002938.XSHE', '003816.XSHE', '300003.XSHE', '300014.XSHE', '300015.XSHE', '300033.XSHE', 
    # '300059.XSHE', '300122.XSHE', '300124.XSHE', '300142.XSHE', '300207.XSHE', '300223.XSHE', '300274.XSHE', '300316.XSHE', '300347.XSHE', '300408.XSHE', '300413.XSHE', '300433.XSHE', '300450.XSHE', '300454.XSHE', '300496.XSHE', 
    # '300498.XSHE', '300529.XSHE', '300595.XSHE', '300601.XSHE', '300628.XSHE', '300661.XSHE', '300750.XSHE', '300751.XSHE', '300759.XSHE', '300760.XSHE', '300763.XSHE', '300782.XSHE', '300866.XSHE', '300896.XSHE', '300919.XSHE', 
    # '300957.XSHE', '300979.XSHE', '300999.XSHE', '600000.XSHG', '600009.XSHG', '600010.XSHG', '600011.XSHG', '600015.XSHG', '600016.XSHG', '600018.XSHG', '600019.XSHG', '600025.XSHG', '600028.XSHG', '600029.XSHG', '600030.XSHG', 
    # '600031.XSHG', '600036.XSHG', '600048.XSHG', '600050.XSHG', '600061.XSHG', '600085.XSHG', '600089.XSHG', '600104.XSHG', '600111.XSHG', '600115.XSHG', '600132.XSHG', '600150.XSHG', '600161.XSHG', '600176.XSHG', '600183.XSHG', 
    # '600188.XSHG', '600196.XSHG', '600219.XSHG', '600276.XSHG', '600309.XSHG', '600332.XSHG', '600346.XSHG', '600352.XSHG', '600362.XSHG', '600383.XSHG', '600406.XSHG', '600426.XSHG', '600436.XSHG', '600438.XSHG', '600460.XSHG', 
    # '600489.XSHG', '600519.XSHG', '600547.XSHG', '600570.XSHG', '600584.XSHG', '600585.XSHG', '600588.XSHG', '600600.XSHG', '600606.XSHG', '600655.XSHG', '600660.XSHG', '600690.XSHG', '600741.XSHG', '600745.XSHG', '600760.XSHG', 
    # '600763.XSHG', '600795.XSHG', '600809.XSHG', '600837.XSHG', '600845.XSHG', '600886.XSHG', '600887.XSHG', '600893.XSHG', '600900.XSHG', '600905.XSHG', '600918.XSHG', '600919.XSHG', '600926.XSHG', '600941.XSHG', '600958.XSHG', 
    # '600989.XSHG', '600999.XSHG', '601006.XSHG', '601009.XSHG', '601012.XSHG', '601021.XSHG', '601066.XSHG', '601088.XSHG', '601100.XSHG', '601111.XSHG', '601117.XSHG', '601138.XSHG', '601155.XSHG', '601166.XSHG', '601169.XSHG', 
    # '601186.XSHG', '601211.XSHG', '601216.XSHG', '601225.XSHG', '601229.XSHG', '601236.XSHG', '601238.XSHG', '601288.XSHG', '601318.XSHG', '601319.XSHG', '601328.XSHG', '601336.XSHG', '601360.XSHG', '601377.XSHG', '601390.XSHG', 
    # '601398.XSHG', '601600.XSHG', '601601.XSHG', '601618.XSHG', '601628.XSHG', '601633.XSHG', '601658.XSHG', '601668.XSHG', '601669.XSHG', '601688.XSHG', '601698.XSHG', '601728.XSHG', '601766.XSHG', '601788.XSHG', '601799.XSHG', 
    # '601800.XSHG', '601808.XSHG', '601816.XSHG', '601818.XSHG', '601825.XSHG', '601838.XSHG', '601857.XSHG', '601865.XSHG', '601868.XSHG', '601877.XSHG', '601878.XSHG', '601881.XSHG', '601888.XSHG', '601898.XSHG', '601899.XSHG', 
    # '601901.XSHG', '601916.XSHG', '601919.XSHG', '601939.XSHG', '601966.XSHG', '601985.XSHG', '601988.XSHG', '601989.XSHG', '601995.XSHG', '601998.XSHG', '603019.XSHG', '603087.XSHG', '603160.XSHG', '603185.XSHG', '603195.XSHG', 
    # '603259.XSHG', '603260.XSHG', '603288.XSHG', '603290.XSHG', '603369.XSHG', '603392.XSHG', '603486.XSHG', '603501.XSHG', '603659.XSHG', '603799.XSHG', '603806.XSHG', '603833.XSHG', '603882.XSHG', '603899.XSHG', '603986.XSHG', 
    # '603993.XSHG', '605499.XSHG', '688005.XSHG', '688008.XSHG', '688012.XSHG', '688036.XSHG', '688065.XSHG', '688111.XSHG', '688126.XSHG', '688169.XSHG', '688363.XSHG', '688396.XSHG', '688561.XSHG', '688599.XSHG', '688981.XSHG']

    from nextt.tasks import group_file_task
    from rqalpha_mod_stu_api.data_source import St_Database_KDataSource


    default_bundle_path = os.path.abspath(os.path.expanduser('D:/.rqalpha/bundle'))
    rqdatac = St_Database_KDataSource(default_bundle_path)

    # 第一步生成买卖信号，按股票分布式处理
    code_list = rqdatac.all_instruments('stock').order_book_id.tolist()

    batch_params = []
    for c in code_list:
      batch_params.append({"benchmark": c})
    start = time.time()
    ret = group_file_task(__file__, batch_params, __config__)
    print(time.time() - start)
    print(json.dumps(ret, indent=2))